import asyncio
import datetime
from io import BytesIO

from app.api.generate.schemas import generate_image_schema
from app.db.session import SessionLocal
from app.models.generate import GenerateImageResult, GenerateImageRecord, GenerateTextRecord, ServerList
from app.utils.curd_base import ModelCRUD
from app.utils.file_storage import save_storage_file, get_storage_file
from app.utils.utils import bs64_to_bytes
from app.utils.websocket_utils import publish_message
from fastapi import UploadFile
from sqlalchemy import or_


class GenerateImageResultCRUD(ModelCRUD):
    model = GenerateImageResult


class GenerateImageRecordCRUD(ModelCRUD):
    model = GenerateImageRecord


def save_generate_image_result(data):
    db = SessionLocal()
    crud = GenerateImageRecordCRUD()
    record_id = data.get('record_id')
    image_base64_list = data.get('image_base64_list', [])
    instance = crud.queryset(db).filter(crud.model.id == record_id, crud.model.state == 1).first()
    if instance is None:
        return
    user_id = instance.user_id
    validate = generate_image_schema.GenerateResultImageCreate.validate
    success = 0
    message = "fail"
    image = None
    image_index = data.get('image_index', 0)
    generate_type = "image"
    try:
        if image_base64_list:

            crud_result_image = GenerateImageResultCRUD()
            print('image_base64_list', len(image_base64_list))
            for image_base64 in image_base64_list:
                image_byte = bs64_to_bytes(image_base64)
                file = UploadFile(file=BytesIO(image_byte), filename="a.png", size=len(image_byte))
                path = save_storage_file(file, "/generate/result/image")
                dit = {"record_id": record_id, "path": path}
                crud_result_image.create(db, data=validate(dit))
                success = 1
                message = "success"
                image = get_storage_file(path)
                message_data = {'image_index': image_index, "record_id": record_id, "success": success, "message": message, "image": image, "generate_type": generate_type}
                asyncio.run(publish_message(str(user_id), message_data, "generate"))

            if image_index+1 >= instance.number:
                instance.state = 2
                db.add(instance)
                db.commit()
                db.refresh(instance)


        else:
            message_data = {'image_index': image_index, "record_id": record_id, "success": success, "message": message, "image": image, "generate_type": generate_type}
            asyncio.run(publish_message(str(user_id), message_data, "generate"))

            instance.state = 3
            db.add(instance)
            db.commit()
            db.refresh(instance)
    except Exception as e:
        print(e)
        message = "逻辑错误"
        message_data = {'image_index': image_index, "record_id": record_id, "success": success, "message": message, "image": image, "generate_type": generate_type}
        asyncio.run(publish_message(str(user_id), message_data, "generate"))

        instance.state = 3
        db.add(instance)
        db.commit()
        db.refresh(instance)


# ------------------------------------------------------------------------------------------------------------------------------------


class GenerateTextRecordCRUD(ModelCRUD):
    model = GenerateTextRecord


def save_generate_text_result(data):
    db = SessionLocal()
    crud = GenerateTextRecordCRUD()
    record_id = data.get('record_id')
    answer = data.get('answer')
    if answer is None:
        return

    instance = crud.queryset(db).filter(crud.model.id == record_id, or_(crud.model.state == 1, crud.model.state == 2)).first()
    if instance is None:
        return
    user_id = instance.user_id

    history = data.get('history')
    success = 0
    message = "fail"
    image_index = data.get('image_index', 0)
    generate_type = "text"
    print("answer len", len(str(answer)))
    try:

        instance.answer = str(answer)
        instance.answer_time = datetime.datetime.now()
        instance.history = history
        instance.state = 2

        db.add(instance)
        db.commit()
        db.refresh(instance)
        success = 1
        message = "success"
        message_data = {'image_index': image_index, "record_id": record_id, "success": success, "message": message, "answer": answer, "generate_type": generate_type, "history": history}
        asyncio.run(publish_message(str(user_id), message_data, "generate"))

    except:
        message = "逻辑错误"
        message_data = {'image_index': image_index, "record_id": record_id, "success": success, "message": message, "answer": answer, "generate_type": generate_type}
        asyncio.run(publish_message(str(user_id), message_data, "generate"))

        instance.state = 3
        db.add(instance)
        db.commit()
        db.refresh(instance)


# ------------------------------------------------------------------------------------------------------------------------------------

class ServerListCRUD(ModelCRUD):
    model = ServerList


def set_server_online(data):
    """
    {"uid": "xxxx", "online": 1}
    :param data:
    :return:
    """
    db = SessionLocal()
    crud = ServerListCRUD()
    instance = crud.queryset(db).filter(crud.model.uid == data.get("uid")).first()
    if instance is None:
        return
    instance.state = data.get("online", 1)
    db.add(instance)
    db.commit()
    db.refresh(instance)

